#!/usr/bin/python

from __future__ import division
import os,threading,re,datetime
from optparse import OptionParser

usage = '''/mnt/ilustre/users/xuefeng.zhan/python/scripts/%prog [options]

---------------------------------------------------------------------------
Sequencing data statistics and quality control program
---------------------------------------------------------------------------
'''

parser = OptionParser(usage = usage)
parser.add_option("-i",dest = "input_file",help = "The sample directories list,each line looks like 'sampleName<TAB>\
sampleDirectory',the sampleDirectory means the directory containing all the gzipped sequencing data")
parser.add_option("-o",dest = "output_directory",default = "QC_out",\
                  help = "The output directory for all samples with quality control files,tables and images [default:%default]")
parser.add_option("-s","--size",dest = "size",default = 151,help = "The length of reads [default:%default]")
(options, args) = parser.parse_args()

def plot(rawdatafalist,size):
    file = open(rawdatafalist)
    try:
        for eachline in file:
            eachline = eachline.split()
            samplename,Fq1,Fq2 = eachline[0],eachline[1],eachline[2]
            stat1,stat2 = samplename + ".1.fq.stat",samplename + ".2.fq.stat"
            threads = []
            t1 = threading.Thread(target = os.system("fastx_quality_stats -i %s -o %s/%s -Q 33"%(Fq1,samplename,stat1)))
            threads.append(t1)
            print "%s start to statistic: %s"%(datetime.datetime.now().strftime('#%Y-%m-%d %H:%M:%S#'),Fq1)
            t2 = threading.Thread(target = os.system("fastx_quality_stats -i %s -o %s/%s -Q 33"%(Fq2,samplename,stat2)))
            threads.append(t2)
            print "%s start to statistic: %s"%(datetime.datetime.now().strftime('#%Y-%m-%d %H:%M:%S#'),Fq2)
            for t in threads:
                t.start()
            for t in threads:
                t.join()
            print "%s statistics complete."%datetime.datetime.now().strftime('#%Y-%m-%d %H:%M:%S#')
            fqcheck_distribute = "perl /mnt/ilustre/users/xuefeng.zhan/scripts/fqcheck_distribute.pl"
            print "%s start to plot,Base Distribution,Base Error Distribution,Quality Scores Distribution."\
            %datetime.datetime.now().strftime('#%Y-%m-%d %H:%M:%S#')
            os.system("%s %s/%s %s/%s -size %i"%(fqcheck_distribute,samplename,stat1,samplename,stat2,size))
        print "%s plot complete."%datetime.datetime.now().strftime('#%Y-%m-%d %H:%M:%S#')
    finally:
        file.close()

def clipper(rawdatafqlist):
    file1 = open(rawdatafqlist)
    file2 = open("trimPairFq.tmp.list","w")
    file3 = open("trimUnpairFq.tmp.list","w")
    file4 = open("rawdata.stat.adapter.xls","w")
    file4.write("sample"+"\t"+"Adapter%"+"\n")
    file4.close()
    try:
        for eachline in file1:
            eachline = eachline.split()
            samplename,Fq1,Fq2 = eachline[0],eachline[1],eachline[2]
            A,B = "AGATCGGAAGAGCACACGTC","AGATCGGAAGAGCGTCGTGT"
            Fq1_clip,Fq2_clip = eachline[1]+".clip.gz",eachline[2]+".clip.gz"
            Fq1_discard,Fq2_discard = eachline[1]+".discard.gz",eachline[2]+".discard.gz"
            print "%s start to reads clip: %s"%(datetime.datetime.now().strftime('#%Y-%m-%d %H:%M:%S#'),samplename)
            os.system("SeqPrep -f %s -r %s -1 %s -2 %s -3 %s -4 %s - q 20 -L 30 -A %s -B %s 2> %s/SeqPrep.log"%(\
                Fq1,Fq2,\
                Fq1_clip,\
                Fq2_clip,\
                Fq1_discard,\
                Fq2_discard,\
                A,B,samplename))
            file = open("%s/SeqPrep.log"%samplename)
            file4 = open("rawdata.stat.adapter.xls","a")
            print "%s start to statistic the adapter rate: %s"%(datetime.datetime.now().strftime('#%Y-%m-%d %H:%M:%S#'),samplename)
            for line in file:
                if re.match("Pairs Processed:",line):
                    pairs = int(line.split("\t")[1])
                if re.match("Pairs With Adapters:",line):
                    adapters = int(line.split("\t")[1])
            file4.write(samplename+"\t"+str(round((adapters*100/pairs),2))+"\n")
            file.close()
            file4.close()
            trim1,trim2,trim_unpaired = samplename+"_trim1.fq",samplename+"_trim2.fq",samplename+"_trim_unpaired.fq"
            os.system("sickle pe -f %s -r %s -o %s/%s -p %s/%s -s %s/%s -t sanger -n -l 30 > %s/sickle.log"%(\
                Fq1_clip,\
                Fq2_clip,\
                samplename,trim1,\
                samplename,trim2,\
                samplename,trim_unpaired,\
                samplename))
            print "%s reads clip complete: %s"%(datetime.datetime.now().strftime('#%Y-%m-%d %H:%M:%S#'),samplename)
            file2.write(samplename+"\t"+os.path.realpath("%s/%s"%(samplename,trim1))+"\t"\
                        +os.path.realpath("%s/%s"%(samplename,trim2))+"\n")
            file3.write(samplename+"\t"+os.path.realpath("%s/%s"%(samplename,trim_unpaired))+"\n")
    finally:
        file1.close()
        file2.close()
        file3.close()
        print "%s All reads clip complete."%datetime.datetime.now().strftime('#%Y-%m-%d %H:%M:%S#')

def stat(rawdatafqlist):
    FastqStat = "/usr/java/latest/bin/java -jar /mnt/ilustre/users/xuefeng.zhan/scripts/FastqStat.jar"
    Dup = "perl /mnt/ilustre/users/xuefeng.zhan/scripts/fastqDup.pl"
    RawDataStatMainFile = "rawdata.stat.main.xls"
    DupFile = "rawdata.stat.dup.xls"
    RawDataStatMainFileaddAdapter = "rawdata.stat.main.xls.addAdapter"
    RawDataStatFile = "rawdata.stat.xls"
    AdapterRateStatFile = "rawdata.stat.adapter.xls"
    print "%s start to create the statistic files."%datetime.datetime.now().strftime('#%Y-%m-%d %H:%M:%S#')
    os.system("%s -i %s > %s"%(FastqStat,rawdatafqlist,RawDataStatMainFile))
    os.system("%s %s > %s"%(Dup,rawdatafqlist,DupFile))
    os.system("perl /mnt/ilustre/app/rna/public_scripts/tabletools_add.pl -i %s -t %s -n 1 -headi T -headt T > %s"\
              %(AdapterRateStatFile,RawDataStatMainFile,RawDataStatMainFileaddAdapter))
    os.system("perl /mnt/ilustre/app/rna/public_scripts/tabletools_add.pl -i %s -t %s -n 1 -headi T -headt T > %s"\
              %(DupFile,RawDataStatMainFileaddAdapter,RawDataStatFile))
    trimPairFqStatFile = "trimdata.pair.stat.xls"
    trimUnpairFqStatFile = "trimdata.unpair.stat.xls"
    trimPairFqListTmp = "trimPairFq.tmp.list"
    trimUnpairFqListTmp = "trimUnpairFq.tmp.list"
    os.system("%s -i %s > %s"%(FastqStat,trimPairFqListTmp,trimPairFqStatFile))
    os.system("%s -i %s > %s"%(FastqStat,trimUnpairFqListTmp,trimUnpairFqStatFile))
    file1 = open(trimPairFqListTmp)
    file2 = open("trimPairFq.list","w")
    for eachline in file1:
        eachline = eachline.split()
        eachline.insert(1,"PE")
        eachline = "\t".join(eachline)+"\n"
        file2.write(eachline)
    file1.close()
    file2.close()
    file1 = open(trimUnpairFqListTmp)
    file2 = open("trimUnpairFq.list","w")
    for eachline in file1:
        eachline = eachline.split()
        eachline.insert(1,"SE")
        eachline = "\t".join(eachline)+"\n"
        file2.write(eachline)
    file1.close()
    file2.close()
    os.remove(trimPairFqListTmp)
    os.remove(trimUnpairFqListTmp)
    os.remove(AdapterRateStatFile)
    os.remove(RawDataStatMainFile)
    os.remove(RawDataStatMainFileaddAdapter)
    os.remove(DupFile)
    os.system("sed 's/#//' rawdata.stat.xls| cut -f 1-3,11-14 > rawdata.show.xls")
    os.system("sed 's/#//' trimdata.pair.stat.xls| cut -f 1-3,11-14 > trimdata.show.xls")
    print "%s files create complete."%datetime.datetime.now().strftime('#%Y-%m-%d %H:%M:%S#')

print "%s start to quality control..."%datetime.datetime.now().strftime('#%Y-%m-%d %H:%M:%S#')
try:
    os.mkdir("%s"%options.output_directory)
except OSError:
    pass
RawDataFqList = "rawdata.fq.list"
file1 = open(options.input_file)
os.chdir("%s"%options.output_directory)
file2 = open(RawDataFqList,"w")
try:
    for eachline in file1:
        eachline = eachline.split()
        samplename = eachline[0]
        pathname = "../"+eachline[1]
        try:
            os.mkdir("%s"%samplename)
        except OSError:
            pass
        files = os.listdir(pathname)
        for eachfile in files:
            unzipfile = eachfile[:-3]
            print "unzip",eachfile
            os.system("gzip -cd %s/%s > ./%s/%s"%(pathname,eachfile,samplename,unzipfile))
        Fq1 = samplename+".1.fq"
        os.rename("%s/%s"%(samplename,"".join(re.findall(".+R1.+","\n".join(files)))[:-3]),"%s/%s"%(samplename,Fq1))
        Fq2 = samplename+".2.fq"
        os.rename("%s/%s"%(samplename,"".join(re.findall(".+R2.+","\n".join(files)))[:-3]),"%s/%s"%(samplename,Fq2))
        file2.write(samplename+"\t"+"%s"%os.path.realpath("%s/%s"%(samplename,Fq1))\
                    +"\t"+"%s"%os.path.realpath("%s/%s"%(samplename,Fq2))+"\n")
finally:
    file1.close()
    file2.close()

threads = []
t1 = threading.Thread(target = plot,args = (RawDataFqList,options.size*2,))
threads.append(t1)
t2 = threading.Thread(target = clipper,args = (RawDataFqList,))
threads.append(t2)
for t in threads:
    t.start()
t.join()
stat(RawDataFqList)
